# A PyQt widget to display 3D scenes in OpenGL, along with interactive camera.
# based on original view3d by Robert Bridson
# Released under the GPLv3

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from PyQt4.QtOpenGL import *
from OpenGL.GL import *
import transform
import math
import numpy
from light import BasicLight

LIGHT_POS = numpy.array([-25.0, 25.0, 25.0])

class View3d(QGLWidget):
    """
    View3d allows basic 3d rendering, taking care of some things for you.
    You can use the default interactive camera, or another camera object.
    Register objects to render, as long as they have initializeGL(),
    initialized(), and paintGL(projViewTransform) methods
    (see axes.py for example), using addItem / removeItem.
    """

    def __init__(self, parent = None):
        QGLWidget.__init__(self, parent)
        self.camera = None
        self.glStatus = False # so we know when we can do GL stuff
        self.items = set() # initially no items to render
        self.light = BasicLight(LIGHT_POS)  
        self.updateTime = 42 # ms
        self.timer = self.startTimer(self.updateTime)
        self.curWaitFrame = 0       # wait for the buffer to be filled
        self.totalWaitFrames = 4
        self.curFrame = 0
        self.totalFrames = 300
        self.saveFrames = False
        self.aspectLocked = False
        
    def timerEvent(self, QTimerEvent):
        if(self.timer == QTimerEvent.timerId()):
            self.update()
            
    def setSaveFrame(self, save):
        self.saveFrames = save
        
    def saveFrame(self):
        if self.curWaitFrame < self.totalWaitFrames:
            self.curWaitFrame += 1
            return

        if self.curFrame < self.totalFrames:
            frameData = QPixmap.grabWindow(self.winId())
            fileNameFormat = str(QDir.currentPath()) + "/export/frame%03d.png"        
            fileName = fileNameFormat % self.curFrame
            if frameData.save(fileName):
                print fileName + " saved" 
            else:
                print "error saving"      
            self.curFrame += 1
        
    def update(self):
        dT = self.updateTime / 1000.0
        for item in self.items:
            item.update(dT)

        self.updateGL()

        # save the first self.totalFrames frames
        if self.saveFrames:
            self.saveFrame()           
        
    def setCamera(self, camera):
        self.camera = camera
        self.update()

    def addItem(self, item):
        self.items.add(item)

        if self.glStatus:
            item.initializeGL()

        self.update()

    def removeItem(self, item):
        self.items.discard(item)
        
    def mousePressEvent(self, event):
        if self.camera != None:
            if self.camera.mousePressEvent(event, self.w, self.h):
                self.update()
   
    def mouseMoveEvent(self, event):
        if self.camera != None and self.camera.mouseMoveEvent(event):
            self.update()

    def mouseReleaseEvent(self, event):
        if self.camera != None and self.camera.mouseReleaseEvent(event):
            self.update()
    
    def initializeGL(self):
        glClearDepth(1)        
        for x in self.items:
            if not x.initialized():
                x.initializeGL()

        self.glStatus = True
        
    def resizeGL(self, w, h):
        if not self.aspectLocked:
            self.w = w
            self.h = h
            self.aspect = self.w / self.h
            self.aspectLocked = True
                        
    def reset(self):
        for item in self.items:
            item.reset()
            
        self.curWaitFrame = 0       # wait for the buffer to be filled
        self.totalWaitFrames = 4
        self.curFrame = 0
        self.totalFrames = 300
        self.saveFrames = False

    def paintGL(self):
        glClearColor(0.9, 0.9, 0.9, 1.0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        mv = self.camera.modelView()
        if self.camera != None:
            pmv = transform.compose(self.camera.projection(), mv)
        else:
            pmv = transform.scaleHiDef(1, self.camera.getAspect(), 1)            
        for x in self.items:
            x.paintGL(mv, pmv, self.light)
